from sqlalchemy import create_engine, Column, Integer, String, DateTime, ForeignKey, Boolean
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import datetime

Base = declarative_base()
engine = create_engine('sqlite:///data/knowledge.db')  # 使用单一数据库文件
Session = sessionmaker(bind=engine)

class ChatHistory(Base):
    __tablename__ = 'chat_history'
    
    id = Column(Integer, primary_key=True)
    role = Column(String)
    content = Column(String)
    timestamp = Column(DateTime, default=datetime.datetime.utcnow)
    
class KnowledgeFile(Base):
    __tablename__ = 'knowledge_files'
    
    id = Column(Integer, primary_key=True)
    path = Column(String, unique=True)  # 文件路径或URL
    file_type = Column(String)  # 文件类型：FILE 或 URL
    status = Column(String)  # 状态：TRAINED 或 PENDING
    added_time = Column(DateTime, default=datetime.datetime.now)
    trained_time = Column(DateTime, nullable=True)

class KnowledgeContent(Base):
    __tablename__ = 'knowledge_contents'
    
    id = Column(Integer, primary_key=True)
    file_id = Column(Integer, ForeignKey('knowledge_files.id'))
    content = Column(String)  # 文件内容
    chunks = Column(String)   # 分块后的内容
    embeddings = Column(String)  # 向量嵌入
    is_complete = Column(Boolean, default=False)  # 添加是否为完整文本的标记
    created_at = Column(DateTime, default=datetime.datetime.now)

# 创建所有表
Base.metadata.create_all(engine) 